﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;

namespace CommonConsole.SM4
{

    public class SM4
    {
        public static readonly int ENCRYPT = 1;
        public static readonly int DECRYPT = 0;
        private static readonly int ROUND = 32;
        private static readonly int BLOCK = 16;

        private static readonly byte[] Sbox = {
                0xD6, 0x90, 0xE9, 0xFE, 0xCC, 0xE1, 0x3D, 0xB7, 0x16, 0xB6, 0x14, 0xC2, 0x28, 0xFB, 0x2C, 0x05,
                0x2B, 0x67, 0x9A, 0x76, 0x2A, 0xBE, 0x04, 0xC3, 0xAA, 0x44, 0x13, 0x26, 0x49, 0x86, 0x06, 0x99,
                0x9C, 0x42, 0x50, 0xF4, 0x91, 0xEF, 0x98, 0x7A, 0x33, 0x54, 0x0B, 0x43, 0xED, 0xCF, 0xAC, 0x62,
                0xE4, 0xB3, 0x1C, 0xA9, 0xC9, 0x08, 0xE8, 0x95, 0x80, 0xDF, 0x94, 0xFA, 0x75, 0x8F, 0x3F, 0xA6,
                0x47, 0x07, 0xA7, 0xFC, 0xF3, 0x73, 0x17, 0xBA, 0x83, 0x59, 0x3C, 0x19, 0xE6, 0x85, 0x4F, 0xA8,
                0x68, 0x6B, 0x81, 0xB2, 0x71, 0x64, 0xDA, 0x8B, 0xF8, 0xEB, 0x0F, 0x4B, 0x70, 0x56, 0x9D, 0x35,
                0x1E, 0x24, 0x0E, 0x5E, 0x63, 0x58, 0xD1, 0xA2, 0x25, 0x22, 0x7C, 0x3B, 0x01, 0x21, 0x78, 0x87,
                0xD4, 0x00, 0x46, 0x57, 0x9F, 0xD3, 0x27, 0x52, 0x4C, 0x36, 0x02, 0xE7, 0xA0, 0xC4, 0xC8, 0x9E,
                0xEA, 0xBF, 0x8A, 0xD2, 0x40, 0xC7, 0x38, 0xB5, 0xA3, 0xF7, 0xF2, 0xCE, 0xF9, 0x61, 0x15, 0xA1,
                0xE0, 0xAE, 0x5D, 0xA4, 0x9B, 0x34, 0x1A, 0x55, 0xAD, 0x93, 0x32, 0x30, 0xF5, 0x8C, 0xB1, 0xE3,
                0x1D, 0xF6, 0xE2, 0x2E, 0x82, 0x66, 0xCA, 0x60, 0xC0, 0x29, 0x23, 0xAB, 0x0D, 0x53, 0x4E, 0x6F,
                0xD5, 0xDB, 0x37, 0x45, 0xDE, 0xFD, 0x8E, 0x2F, 0x03, 0xFF, 0x6A, 0x72, 0x6D, 0x6C, 0x5B, 0x51,
                0x8D, 0x1B, 0xAF, 0x92, 0xBB, 0xDD, 0xBC, 0x7F, 0x11, 0xD9, 0x5C, 0x41, 0x1F, 0x10, 0x5A, 0xD8,
                0x0A, 0xC1, 0x31, 0x88, 0xA5, 0xCD, 0x7B, 0xBD, 0x2D, 0x74, 0xD0, 0x12, 0xB8, 0xE5, 0xB4, 0xB0,
                0x89, 0x69, 0x97, 0x4A, 0x0C, 0x96, 0x77, 0x7E, 0x65, 0xB9, 0xF1, 0x09, 0xC5, 0x6E, 0xC6, 0x84,
                0x18, 0xF0, 0x7D, 0xEC, 0x3A, 0xDC, 0x4D, 0x20, 0x79, 0xEE, 0x5F, 0x3E, 0xD7, 0xCB, 0x39, 0x48
            };

        private static readonly uint[] CK = {
                0x00070e15, 0x1c232a31, 0x383f464d, 0x545b6269, 0x70777e85, 0x8c939aa1, 0xa8afb6bd, 0xc4cbd2d9,
                0xe0e7eef5, 0xfc030a11, 0x181f262d, 0x343b4249, 0x50575e65, 0x6c737a81, 0x888f969d, 0xa4abb2b9,
                0xc0c7ced5, 0xdce3eaf1, 0xf8ff060d, 0x141b2229, 0x30373e45, 0x4c535a61, 0x686f767d, 0x848b9299,
                0xa0a7aeb5, 0xbcc3cad1, 0xd8dfe6ed, 0xf4fb0209, 0x10171e25, 0x2c333a41, 0x484f565d, 0x646b7279
            };

        private static readonly uint[] FK = { 0xa3b1bac6, 0x56aa3350, 0x677d9197, 0xb27022dc };

        private static uint Rotl(uint a, int b)
        {
            return (((a) << (b)) | ((a) >> (32 - (b))));
        }

        private static uint ByteSub(uint A)
        {
            return (uint)(((Sbox[(A >> 24) & 0xFF] & 0xFF) << 24) | (Sbox[A >> 16 & 0xFF] & 0xFF) << 16
                    | ((Sbox[(A >> 8) & 0xFF] & 0xFF) << 8) | (Sbox[A & 0xFF] & 0xFF));
        }

        private static uint L1(uint B)
        {
            return B ^ Rotl(B, 2) ^ Rotl(B, 10) ^ Rotl(B, 18) ^ Rotl(B, 24);
        }

        private static uint L2(uint B)
        {
            return B ^ Rotl(B, 13) ^ Rotl(B, 23);
        }

        private static void SMS4Crypt(byte[] Input, byte[] Output, uint[] rk)
        {
            uint r;
            uint mid;
            uint[] x = new uint[4];
            uint[] tmp = new uint[4];
            for (int i = 0; i < 4; i++)
            {
                tmp[0] = (uint)(Input[0 + 4 * i] & 0xff);
                tmp[1] = (uint)(Input[1 + 4 * i] & 0xff);
                tmp[2] = (uint)(Input[2 + 4 * i] & 0xff);
                tmp[3] = (uint)(Input[3 + 4 * i] & 0xff);
                x[i] = (uint)((tmp[0] << 24) | (tmp[1] << 16) | (tmp[2] << 8) | tmp[3]);
            }
            for (r = 0; r < 32; r += 4)
            {
                mid = (uint)(x[1] ^ x[2] ^ x[3] ^ rk[r + 0]);
                mid = ByteSub(mid);
                x[0] = x[0] ^ L1(mid);

                mid = (uint)(x[2] ^ x[3] ^ x[0] ^ rk[r + 1]);
                mid = ByteSub(mid);
                x[1] = x[1] ^ L1(mid);

                mid = (uint)(x[3] ^ x[0] ^ x[1] ^ rk[r + 2]);
                mid = ByteSub(mid);
                x[2] = x[2] ^ L1(mid);

                mid = (uint)(x[0] ^ x[1] ^ x[2] ^ rk[r + 3]);
                mid = ByteSub(mid);
                x[3] = x[3] ^ L1(mid);
            }

            // Reverse
            for (int j = 0; j < 16; j += 4)
            {
                Output[j] = (byte)(x[3 - j / 4] >> 24 & 0xFF);
                Output[j + 1] = (byte)(x[3 - j / 4] >> 16 & 0xFF);
                Output[j + 2] = (byte)(x[3 - j / 4] >> 8 & 0xFF);
                Output[j + 3] = (byte)(x[3 - j / 4] & 0xFF);
            }
        }

        private static void SMS4KeyExt(byte[] Key, uint[] rk, int CryptFlag)
        {
            uint r, mid;
            uint[] x = new uint[4];
            uint[] tmp = new uint[4];
            for (int i = 0; i < 4; i++)
            {
                tmp[0] = (uint)(Key[0 + 4 * i] & 0xFF);
                tmp[1] = (uint)(Key[1 + 4 * i] & 0xff);
                tmp[2] = (uint)(Key[2 + 4 * i] & 0xff);
                tmp[3] = (uint)(Key[3 + 4 * i] & 0xff);
                x[i] = tmp[0] << 24 | tmp[1] << 16 | tmp[2] << 8 | tmp[3];
            }
            x[0] ^= FK[0];
            x[1] ^= FK[1];
            x[2] ^= FK[2];
            x[3] ^= FK[3];
            for (r = 0; r < 32; r += 4)
            {
                mid = x[1] ^ x[2] ^ x[3] ^ CK[r + 0];
                mid = ByteSub(mid);
                rk[r + 0] = x[0] ^= L2(mid); // rk0=K4

                mid = x[2] ^ x[3] ^ x[0] ^ CK[r + 1];
                mid = ByteSub(mid);
                rk[r + 1] = x[1] ^= L2(mid); // rk1=K5

                mid = x[3] ^ x[0] ^ x[1] ^ CK[r + 2];
                mid = ByteSub(mid);
                rk[r + 2] = x[2] ^= L2(mid); // rk2=K6

                mid = x[0] ^ x[1] ^ x[2] ^ CK[r + 3];
                mid = ByteSub(mid);
                rk[r + 3] = x[3] ^= L2(mid); // rk3=K7
            }

            // ½âÃÜÊ±ÂÖÃÜÔ¿Ê¹ÓÃË³Ðò£ºrk31,rk30,...,rk0
            if (CryptFlag == DECRYPT)
            {
                for (r = 0; r < 16; r++)
                {
                    mid = rk[r];
                    rk[r] = rk[31 - r];
                    rk[31 - r] = mid;
                }
            }
        }

        private static void xors(byte[] iv, byte[] Input)
        {
            for (int i = 0; i < Input.Length; i++)
            {
                Input[i] = (byte)(Input[i] ^ iv[i]);
            }
        }

        public static void SMS4_ECB(byte[] Input, int inLen, byte[] key, byte[] Output, int CryptFlag)
        {
            int point = 0;
            uint[] round_key = new uint[ROUND];
            SMS4KeyExt(key, round_key, CryptFlag);
            byte[] input = new byte[16];
            byte[] output = new byte[16];

            while (inLen >= BLOCK)
            {
                Array.Copy(Input, point, input, 0, 16);
                SMS4Crypt(input, output, round_key);
                Array.Copy(output, 0, Output, point, BLOCK);
                inLen -= BLOCK;
                point += BLOCK;
            }
        }

        public static void SMS4_BCB(byte[] Input, int inLen, byte[] key, byte[] iv, byte[] Output, int CryptFlag)
        {
            int point = 0;
            uint[] round_key = new uint[ROUND];
            SMS4KeyExt(key, round_key, CryptFlag);
            byte[] input = new byte[16];
            byte[] output = new byte[16];
            byte[] tmpIV = new byte[16];

            if (CryptFlag == ENCRYPT)
            {
                Array.Copy(iv, 0, tmpIV, 0, 16);

                while (inLen >= BLOCK)
                {
                    Array.Copy(Input, point, input, 0, 16);
                    xors(tmpIV, input);
                    SMS4Crypt(input, output, round_key);
                    Array.Copy(output, 0, tmpIV, 0, 16);
                    Array.Copy(output, 0, Output, point, BLOCK);
                    inLen -= BLOCK;
                    point += BLOCK;
                }
            }
            else
            {
                Array.Copy(iv, 0, tmpIV, 0, 16);

                while (inLen >= BLOCK)
                {
                    Array.Copy(Input, point, input, 0, 16);
                    SMS4Crypt(input, output, round_key);
                    xors(tmpIV, output);
                    Array.Copy(input, 0, tmpIV, 0, 16);
                    Array.Copy(output, 0, Output, point, BLOCK);
                    inLen -= BLOCK;
                    point += BLOCK;
                }
            }
        }

        /// <summary>
        /// 长度不为16的整数倍时，填充0xff，并再增加16字节，用于说明数据长度
        /// </summary>
        /// <param name="Input"></param>
        /// <param name="inLen"></param>
        /// <param name="key"></param>
        /// <param name="Output"></param>
        /// <param name="CryptFlag"></param>
        public static void ECB_Fill(byte[] Input, int inLen, byte[] key, System.IO.Stream Output, int CryptFlag)
        {
            using (System.IO.MemoryStream inputStream = new System.IO.MemoryStream(Input, 0, inLen))
            {
                ECB_Fill(inputStream, key, Output, CryptFlag);
            }
        }


        private static void ECB_FILL_Encrypt(System.IO.Stream Input, byte[] key, System.IO.Stream Output)
        {
            uint[] round_key = new uint[ROUND];
            SMS4KeyExt(key, round_key, ENCRYPT);
            byte[] inputBuf = new byte[BLOCK];
            byte[] outputBuf = new byte[BLOCK];
          
            int len = 0;
            uint inLen = 0;
            do
            {
                len = Input.Read(inputBuf, 0, BLOCK);
                inLen += (uint)len;
                if (len < BLOCK)
                {
                    for (var i = len; i < BLOCK; i++)
                    {
                        inputBuf[i] = 0xFF;
                    }

                }


                SMS4Crypt(inputBuf, outputBuf, round_key);

                Output.Write(outputBuf, 0, outputBuf.Length);
            }
            while (len == BLOCK);



            //加密时把数据长度写入结尾
            Array.Clear(inputBuf, 0, inputBuf.Length);
            //前15个字节填充为FILLECB_FILLLEN
            var fillBytes = System.Text.Encoding.UTF8.GetBytes("FILLECB_FILLLEN");
            Array.Copy(fillBytes, inputBuf, 15);

            //最后一个字节填充 长度
            inputBuf[15] = (byte)(BLOCK - len); ;

            SMS4Crypt(inputBuf, outputBuf, round_key);

            Output.Write(outputBuf, 0, outputBuf.Length);
        }

        private static void ECB_FILL_Decrypt(System.IO.Stream Input, byte[] key, System.IO.Stream Output)
        {
            uint[] round_key = new uint[ROUND];
            SMS4KeyExt(key, round_key, DECRYPT);
            byte[] inputBuf = new byte[BLOCK];
            byte[] outputBuf = new byte[BLOCK];
            //输入流不一定支持获取长度、输出流可能不支持重设长度，需要建立两个组的缓冲区用于截掉填充数据
            int frameCount = 0;
            byte[] fillerBuf1 = new byte[BLOCK];
            byte[] fillerBuf2 = new byte[BLOCK];

            int len = 0;
            do
            {
                len = Input.Read(inputBuf, 0, BLOCK);

                if (len < BLOCK)
                {
                    break;
                }

                frameCount++;

                if (frameCount > 1)
                {
                    if (frameCount > 2)
                        Output.Write(fillerBuf1, 0, fillerBuf1.Length);

                    Array.Copy(fillerBuf2, fillerBuf1, BLOCK);
                }

                SMS4Crypt(inputBuf, fillerBuf2, round_key);
                
            }
            while (len == BLOCK);



            //解密时，把末尾的填充数据截掉
            if (frameCount >= 2)
            { 
                //判断标识
                var v = System.Text.Encoding.UTF8.GetString(fillerBuf2, 0, 15);

                if (v == "FILLECB_FILLLEN")
                {
                    var fillLen = fillerBuf2[15];

                    //如果fillLen 大于 16，则数据错误
                    if (fillLen <= BLOCK)
                        Output.Write(fillerBuf1, 0, BLOCK - fillLen);
                    else
                    {
                        Output.Write(fillerBuf1, 0, fillerBuf1.Length);
                        Output.Write(fillerBuf2, 0, fillerBuf2.Length);
                    }
                }
                else
                {
                    Output.Write(fillerBuf1, 0, fillerBuf1.Length);
                    Output.Write(fillerBuf2, 0, fillerBuf2.Length);
                }
            }
            else if (frameCount == 1)
            {
                Output.Write(fillerBuf2, 0, fillerBuf2.Length);
            }
        }

        public static void ECB_Fill(System.IO.Stream Input, byte[] key, System.IO.Stream Output, int CryptFlag)
        {
            if (CryptFlag == ENCRYPT)
                ECB_FILL_Encrypt(Input,key,Output);
            else
                ECB_FILL_Decrypt(Input, key, Output);            
        }



        public static void Test()
        {
            string v = "hello world";

            var bytes = System.Text.Encoding.UTF8.GetBytes(v);

            var key =System.Text.Encoding.UTF8.GetBytes( "1234567812345678");
            //带填充加密，数据可为任意长度
            using (var ensm = new System.IO.MemoryStream())
            using (var desm = new System.IO.MemoryStream())
            using (var sr = new System.IO.StreamReader(desm))
            {
                ECB_Fill(bytes, bytes.Length, key, ensm, ENCRYPT);

                ensm.Position = 0;

                ECB_Fill(ensm, key,desm, DECRYPT);

                desm.Position = 0;

                Console.WriteLine(sr.ReadToEnd());
            }


            v = "12345678123456781234567812345678";
            bytes = System.Text.Encoding.UTF8.GetBytes(v);
            var obytes = new byte[32];

            //不带填充加密，数据长度必需是16的整数倍
            SMS4_ECB(bytes, bytes.Length, key, obytes, ENCRYPT);

            //解析不带填充的数据
            using (var ensm = new System.IO.MemoryStream(obytes))
            using (var desm = new System.IO.MemoryStream())
            using (var sr = new System.IO.StreamReader(desm))
            {              
                ensm.Position = 0;

                ECB_Fill(ensm, key, desm, DECRYPT);

                desm.Position = 0;

                Console.WriteLine(sr.ReadToEnd());
            }

  
            
        }
    }
}
